package com.hzdaba.aiapi.handlerInterceptor;

import com.hzdaba.aiapi.properties.IpKeyProperties;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.HandlerInterceptor;
import org.springframework.web.servlet.ModelAndView;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.TimeUnit;

/**
 * 限流
 *
 * @author 刘逸晖
 */
@Component
public class CurrentLimitingHandlerInterceptor implements HandlerInterceptor {

    /**
     * 访问 Redis
     */
    @Autowired
    private RedisTemplate<String, Object> redisTemplate;

    /**
     * ip key 相关配置
     */
    @Autowired
    private IpKeyProperties ipKeyProperties;

    /**
     * 一段时间内的最大访问次数
     */
    @Value("${aiapi.request.remoteaddr.limiting.maxNumber}")
    private int maxNumber;

    /**
     * 请求中存放错误信息的 key
     */
    @Value("${aiapi.request.attribute.error.key}")
    private String keyOfError;

    /**
     * 日志记录器
     */
    private static final Logger logger = LoggerFactory.getLogger(CurrentLimitingHandlerInterceptor.class);

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        //获取远程地址
        String remoteAddr = request.getRemoteAddr();

        //检查 Redis 中是否有缓存这个 ip
        Integer count = (Integer) this.redisTemplate.opsForHash().get(this.ipKeyProperties.getPrefix(), remoteAddr);

        if (count != null && count > this.maxNumber) {//该客户端访问的次数过多，拒绝访问
            response.setStatus(403);
            request.setAttribute(this.keyOfError, "抱歉，您的访问次数过多，请稍后再试");
            request.getRequestDispatcher("/error/error").forward(request, response);

//            CurrentLimitingHandlerInterceptor.logger.info("客户端访问次数过多，不与访问");
            return false;
        } else if (count != null) {//已经有了技术，+1
            this.redisTemplate.opsForHash().increment(this.ipKeyProperties.getPrefix(), remoteAddr, 1);
        } else {
            this.redisTemplate.opsForHash().put(this.ipKeyProperties.getPrefix(), remoteAddr, 1);
        }

        //每次请求都设置 key 的过期时间，好像没必要，应该只有Redis 启动后，每个 ip 第一次访问需要设置，可能弄一个定时任务定时删除更合适吧

        if (ThreadLocalRandom.current().nextInt(100) == 0) {//随机数是0，要设置过期时间
            this.redisTemplate.expire(this.ipKeyProperties.getPrefix(), this.ipKeyProperties.getExpiry(), TimeUnit.MINUTES);
        }

        return true;
    }

    @Override
    public void postHandle(HttpServletRequest request, HttpServletResponse response, Object handler, ModelAndView modelAndView) throws Exception {
        HandlerInterceptor.super.postHandle(request, response, handler, modelAndView);
    }

    @Override
    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex) throws Exception {
        HandlerInterceptor.super.afterCompletion(request, response, handler, ex);
    }
}
